import asyncio
import websockets
from loguru import logger
from gpt import chatbot

# 处理WebSocket连接的回调函数
async def handle_websocket(websocket, path):
    try:
        while True:
            # 接收客户端发送的消息
            message = await websocket.recv()
            logger.debug(f"Received message: {message}")
         
            # 发送消息给客户端
            msg = message.split(":")[1]
            # await websocket.send("我是AI助手GPT")
            for data in chatbot.ask(f"{msg},请尽量简短回复"):
                response = data["message"]
            logger.debug(response)
            await websocket.send(response)

            if msg == "exit":
                logger.info('退出server')
                # 收到 "quit" 消息后退出循环并关闭连接
                break

    except websockets.exceptions.ConnectionClosedError:
        # print("WebSocket connection closed.")
        pass
    # 关闭WebSocket连接
    await websocket.close()


logger.info('创建WebSocket服务器')
start_server = websockets.serve(handle_websocket, '0.0.0.0', 8765)
# 创建事件循环
loop = asyncio.get_event_loop()
# 启动事件循环和WebSocket服务器
server = loop.run_until_complete(start_server)

try:
    # 运行事件循环
    logger.info("start")
    loop.run_forever()
finally:
    # 关闭WebSocket服务器
    logger.info("stop")
    server.close()
    loop.run_until_complete(server.wait_closed())

    # 关闭事件循环
    loop.close()